#include "precompiled.h"

#include <gsl/gsl_fft_complex.h>
#include <gsl/gsl_complex_math.h>

#include "Config.h"
#include "Util.h"
#include "Environment.h"
#include "Robot.h"
#include "ViacarTeam2.h"

using namespace boost::numeric::ublas;

#define LEFT -1
#define RIGHT 1

// filter for use in downsampling. Cutoff is at 25Hz
#define H2_L 2
static const double H2_SOS[] = {
1.0000000, -1.8857977, 1.0000000, 1.0000000, -1.9253500, 0.9337392,
1.0000000, -1.4150224, 1.0000000, 1.0000000, -1.8295524, 0.8380458};
static const double H2_G[] = {0.0734597, 0.0145192};

#define H3_L 1
static const double H3_SOS[] = {
1.000000000000000, 2.000000000000000, 1.000000000000000, 1.000000000000000, -1.991114292201654, 0.991153595868935};
static const double H3_G[] = {0.000009825916820};

ViacarTeam2::ViacarTeam2() :
	MaxSpeed (500.0),
	MaxWheelAngle (0.34906585),
	Width (15.0),
	Length (25.0),
	CenterOfMass (12.0),
	Kp (100.0),
	Ki (0),
	Kd (0),
	IntegralPoints (10),
	SensorBits(10),
	SatThresh(20.0),
	CrossingFwdThresh(100.0),
	CrossingErrorThresh(100.0),
	LostThresh(100.0)
{
	typedef json::grammar<char>::array jarray;
	typedef json::grammar<char>::object jobject;

	jobject &o = Config::Get("Robot");

	Config::TryAssign(MaxSpeed, "MaxSpeed", o);
	Config::TryAssign(MaxWheelAngle, "MaxWheelAngle", o);
	Config::TryAssign(Width, "Width", o);
	Config::TryAssign(Length, "Length", o);
	Config::TryAssign(CenterOfMass, "CenterOfMass", o);
	Config::TryAssign(Kp, "Kp", o);
	Config::TryAssign(Ki, "Ki", o);
	Config::TryAssign(Kd, "Kd", o);
	Config::TryAssign(IntegralPoints, "IntegralPoints", o);
	Config::TryAssign(SensorBits, "SensorBits", o);
	Config::TryAssign(SatThresh, "SatThresh", o);
	Config::TryAssign(CrossingFwdThresh, "CrossingFwdThresh", o);
	Config::TryAssign(CrossingErrorThresh, "CrossingErrorThresh", o);
	Config::TryAssign(LostThresh, "LostThresh", o);

	jobject &a = boost::any_cast<jobject>(*o.find("Sensors")->second);
	Sensors.resize(a.size());
	unsigned int i = 0;
	for(jobject::const_iterator it = a.begin(); i < a.size(); it++, i++)
	{
		jobject &s = boost::any_cast<jobject>(*it->second);
		SensorConstants &c = Sensors[i];

		mSensorMap[it->first] = i;

		jarray &xyz = boost::any_cast<jarray>(*s.find("Xyz")->second);
		jarray &hpr = boost::any_cast<jarray>(*s.find("Hpr")->second);

		c.Name = it->first;
		Config::AssignArray(c.Xyz, xyz);
		Config::AssignArray(c.Hpr, hpr);
		Config::TryAssign(c.Size, "Size", s);
		Config::TryAssign(c.Turns, "Turns", s);
		Config::TryAssign(c.Gain, "Gain", s);
	}
	mSensors.resize(Sensors.size());

	iir_sos_init(&mH3, H3_SOS, H3_G, H3_L);

	// set up the monte carlo integration and random number generator
	mWorkspace = gsl_monte_plain_alloc(2);
	gsl_monte_plain_init(mWorkspace);
	mRng = gsl_rng_alloc(gsl_rng_taus);
}

ViacarTeam2::~ViacarTeam2() {
	gsl_monte_plain_free(mWorkspace);
	gsl_rng_free(mRng);
}

// sets the mSpeed by clamping within range
void ViacarTeam2::SetSpeed(double value) {
	// TODO: introduce time delay
	mSpeed = Util::clamp<double>(value, 0, 1.0);
}

// sets the wheel angle by claming within max range
void ViacarTeam2::SetWheelAngle(double value) {
	if(mCanUpdateDirection) {
		mDirection = Util::clamp(value, -1.0, 1.0);
	}
}

void ViacarTeam2::InitSensor(SensorData &d, const SensorConstants &c) {
	d.monte.m1.resize(4, 4);
	Util::MakeEulerMatInPlace(d.monte.m1, c.Hpr[0], c.Hpr[1], c.Hpr[2]);
	d.monte.m1(0, 3) = c.Xyz[0];
	d.monte.m1(1, 3) = c.Xyz[1];
	d.monte.m1(2, 3) = c.Xyz[2];
	d.monte.env = mEnv;

	iir_sos_init(&d.h2, H2_SOS, H2_G, H2_L);

	mEnv->GetResult().AddSignal(c.Name);
	mEnv->GetResult().AddSignal(c.Name + "_inv");
}

void ViacarTeam2::Start() {
	for(int i = 0; i < mSensors.size(); i++) {
		InitSensor(mSensors[i], Sensors[i]);
	}
	mEnv->GetResult().AddSignal("error");
	mEnv->GetResult().AddSignal("fwd_error");
	mEnv->GetResult().AddSignal("wheelUpdate");
	mEnv->GetResult().AddSignal("sideOfTrack");
	mEnv->GetResult().AddSignal("frontCenter");
	mEnv->GetResult().AddSignal("speed");
	mEnv->GetResult().AddSignal("direction");
	mEnv->GetResult().AddSignal("CrossingFwdThresh");
	mEnv->GetResult().AddSignal("CrossingErrorThresh");


	// the robot is now allow to ask the environment questions,
	// like the magnetic field strength at certain points

	SetSpeed(0.0);		// set full speed ahead
	mDirection = 0;
	mPreviousError = 0;
	mIntegral = 0;
	mLastX = 0;
	mCanUpdateDirection = false;
	mEquilibriated = false;
	iir_sos_reset(&mH3);
	mKp = Kp;
	mSideOfTrack = 0;
}

// add an offset and multiply by gain
double ViacarTeam2::DoSensorAnalog(double emf, const SensorConstants &c) {
	double y = emf * c.Gain + 1.5;
	return y;
}

void ViacarTeam2::DoSensor(SensorData &d, const SensorConstants &c) {
	// block dc, rectify, and filter
	d.sensed = Sense(d, c);
	d.sampled = Sample(DoSensorAnalog(d.sensed, c));
	d.blocked = d.sampled - 512;
	d.rectified = abs(d.blocked);
	d.filtered = iir_sos_filter(&d.h2, d.rectified);			// low pass filter
	if(d.filtered < 1.0) {
		d.filtered = 1.0;
	}
}

//Every step, we are allowed to update our speed and direction, that's it.
void ViacarTeam2::Step(double t, double dt) {
	// are we allowed to update our direction on this step?
	int x = (int)(t/0.02);
	mCanUpdateDirection = false;
	mWheelUpdate = 0;
	if(x != mLastX) {
		mLastX = x;
		mCanUpdateDirection = true;
		mWheelUpdate = 1;
	}

	int l = mSensorMap["left"];
	int r = mSensorMap["right"];
	int fl = mSensorMap["front_left"];
	int fr = mSensorMap["front_right"];
	int fc = mSensorMap["front_center"];
	int fw = mSensorMap["front_fwd"];

	// process the signal coming in on all sensors
	for(int i = 0; i < mSensors.size(); i++) {
		DoSensor(mSensors[i], Sensors[i]);
	}

	// run the front center sensor through a steeper filter
	mFrontCenter = iir_sos_filter(&mH3, mSensors[fc].rectified);

	mError = 1.0/mSensors[r].filtered - 1.0/mSensors[l].filtered;
	mFwdError = 1.0/mSensors[fr].filtered - 1.0/mSensors[fl].filtered;

	// for the first 50ms, just sit and wait for the sensor values to equilibriate
	if(t < 0.05) {
		return;
	} else if(!mEquilibriated) {
		// grab the value of the sensor
		for(int i = 0; i < mSensors.size(); i++) {
			mSensors[i].centered = mSensors[i].filtered;
			std::cerr << "centered[" << i << "] = " << mSensors[i].centered;
		}
		SetSpeed(1.0);
		mEquilibriated = true;
	}

	// determine which side of the track we are on
	if(mSensors[l].filtered < mSensors[l].centered && mSensors[r].filtered > mSensors[r].centered) {
		mSideOfTrack = LEFT;
	} else if(mSensors[r].filtered < mSensors[r].centered && mSensors[l].filtered > mSensors[l].centered) {
		mSideOfTrack = RIGHT;
	}
	// otherwise indeterminate. Keep value the same

	/*
	// determine which scenario we are in
	// if both left and right sensors are below a threshold, then we are lost
	if(mSensors[l].filtered < LostThresh && mSensors[r].filtered < LostThresh) {
		mMode = LOST;
	} else if(mSensors[fw].filtered > CrossingFwdThresh && mFwdError > CrossingErrorThresh) {
		// determine if we are at a crossing. We are at a crossing if the forward sensor
		// exceeds a threshold and the error exceeds a threshold
		mMode = CROSSING;
	} else if(mSensors[fc].filtered < 0.25 * mSensors[fc].centered) {
		// determine if we are on a turn
	} else {
		mMode = LINEAR;
	}
	*/

	// should we slow down?
	/*
	if(mSensors[fc].filtered < 0.25 * mSensors[fc].centered) {
		mKp = 3.0*Kp;
		SetSpeed(0.5);
	} else if(mFrontCenter > 0.75 * mSensors[fc].centered) {
		mKp = Kp;
		SetSpeed(1.0);
	}
	*/

	if(mCanUpdateDirection) {
		// determine if we are in linear or saturation mode
		//if(mSideOfTrack == LEFT && mSensors[l].filtered < SatThresh) {
			// we are left saturated. turn hard right
		//	SetWheelAngle(-1.0);
		//} else if(mSideOfTrack == RIGHT && mSensors[r].filtered < SatThresh) {
		//	SetWheelAngle(1.0);
		//} else if(mSensors[fw].filtered > 100) {
			// we're at a crossing. Stay the course.
		//} else {
			// linear mode
			// set the speed inversely to the forward error
			//SetSpeed(1.0 - 10.0 * abs(mFwdError));
			
			// Increase kp as speed decreases
			//const double s = 3.0;
			//double kp = Kp * (s * (1.0 - mSpeed) + 1.0);
			//double kp = Kp;
			double error = mFwdError;
			mIntegral += (error * 0.02);
			double derivative = (error - mPreviousError) / 0.02;
			mOutput = (mKp * error) + (Ki * mIntegral) + (Kd * derivative);
			mPreviousError = error;

			SetWheelAngle(mOutput);
			
		//}
	}
}

// callback for use in the monte carlo integrator
// x_array is an array of size 2 that is the point we should evaluate the integral at
// dim will always be 2
// params is of type my_monte_params
// the function we want to evaluate is (B dot n), where B is the magnetic field vector
// and n is the normal vector to the sensor
static double __monte_carlo_func(double *x_array, size_t dim, void *vpParams) {
	typedef vector<double> vec;

	ViacarTeam2::my_monte_params *params = static_cast<ViacarTeam2::my_monte_params *>(vpParams);

	vec p0(4);
	p0[0] = x_array[0];
	p0[1] = x_array[1];
	p0[2] = 0;
	p0[3] = 1.0;
	
	// transform the point to world coordinates
	vec p = prod(params->m2, p0);
	//std::cerr << "p0: " << p0 << ", p: " << p << "\n";

	// get the magnetic field and dot it with the normal vector
	vec B = params->env->GetTrack().GetFieldAt(project(p, range(0, 3)));
	assert(B[0] == 0 && B[1] == 0 && B[2] == 0);
	double val = inner_prod(B, params->n);
	//std::cerr << val << "\n";
	return val;
}

// gets the induced voltage at the output of the coils.
double ViacarTeam2::Sense(SensorData &d, const SensorConstants &c) {
	double flux, abs_error;

	// untransformed normal vector (pointing down positive z axis)
	vec n(4);
	n[0] = 0;
	n[1] = 0;
	n[2] = 1.0;
	n[3] = 1.0;

	d.monte.m2 = prod(mEnv->GetCarMat(), d.monte.m1);
	//std::cerr << "m2: " << params.m2 << "\n";

	// see http://fly.srk.fer.hr/~unreal/theredbook/appendixg.html for transforming normal vectors
	// compute as n2 = transpose(invert(M)) * n1;
	matrix<double> m2Inverse(4, 4);
	bool ok = InvertMatrix(d.monte.m2, m2Inverse); 
	assert(ok);
	matrix<double> m2InverseTranspose = trans(m2Inverse);
	d.monte.n = project(prod(m2InverseTranspose, n), range(0, 3));
	//std::cerr << "n: " << params.n << "\n";

	//vec v1(4);
	//v1[0] = 0;
	//v1[1] = 0;
	//v1[2] = 0;
	//v1[3] = 1.0;
	//vec v2 = prod(params.m2, v1);
	//std::cerr << "v2: " << v2 << "\n";

	double xl[] = {-c.Size/2, -c.Size/2};
	double xu[] = {c.Size/2, c.Size/2};
	gsl_monte_function func = {__monte_carlo_func, 2, &d.monte};
	gsl_monte_plain_integrate (&func, xl, xu, 2, IntegralPoints, mRng, mWorkspace, &flux, &abs_error);
	
	// the magnitude of the induced voltage is the time derivative of the flux
	// need to adjust for units. Integration area was in cm, but field was in Tesla
	double emf = 2.0 * M_PI * c.Turns * mTrack->Freq * flux * cos(2 * M_PI * mTrack->Freq * mEnv->T());
	return emf;
}

unsigned int ViacarTeam2::Sample(double v) {
	const int max = 1 << SensorBits;
	double y = v * max / 3.0;
	// quantize and clamp
	int quant = Util::clamp((int)y, 0, max);

	return (unsigned int)quant;
}

void ViacarTeam2::Save(SimResult::Data &data) {
	SimResult &res = mEnv->GetResult();
	for(int i = 0; i < mSensors.size(); i++) {
		res.SaveSignal(Sensors[i].Name, mSensors[i].filtered);
		res.SaveSignal(Sensors[i].Name + "_inv", 1.0/mSensors[i].filtered);
	}
	res.SaveSignal("error", mError);
	res.SaveSignal("fwd_error", mFwdError);
	res.SaveSignal("wheelUpdate", mWheelUpdate);
	res.SaveSignal("sideOfTrack", mSideOfTrack);
	res.SaveSignal("frontCenter", mFrontCenter);
	res.SaveSignal("speed", mSpeed);
	res.SaveSignal("direction", mDirection);
}
